from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from app.models.order import Order   # 新增
from app.services.broker import dry_run   # 新增
from app.services.market_crawler import crawl_1m, CH, crawl_all_1m
from app.strategy.do_t import do_t_signal
from app.services.twap import twap_worker
from app.services.risk import update_nav
import asyncio
import json
import pandas as pd
from loguru import logger
# 全局假持仓（代码 -> 股数）
FAKE_PORT = {'sz000001': 5000, 'sz000002': 3000, 'sh600519': 1000}

router = APIRouter(prefix="/api/market")

@router.get("/crawl-all")
def crawl_all():
    crawl_all_1m()
    return {"msg": "crawled all stocks"}

from app.gnn.infer import gnn_signal
async def latest_tick():
    while True:
        crawl_all_1m()                     # 轮询全部标的
        df = CH.query_dataframe('''
                                SELECT code, close AS price, dt
                                FROM tick_1m
                                WHERE dt >= now() - INTERVAL 1 MINUTE
                                ''')
        print("df:",df)
        if df.empty:          # 空数据用假行情
            df = pd.DataFrame([
                {'code': 'sz000001', 'price': 14.5, 'dt': pd.Timestamp.now()},
                {'code': 'sz000002', 'price': 28.3, 'dt': pd.Timestamp.now()},
                {'code': 'sh600519', 'price': 1720, 'dt': pd.Timestamp.now()},
            ])
        ticks = df.to_dict(orient='records')
        
        total_mv = sum(tick['price'] * FAKE_PORT.get(tick['code'], 0) for tick in ticks)
        update_nav(total_mv)

        gnn = gnn_signal()
        for _, row in df.iterrows():
            tick = row.to_dict()
            for key, value in tick.items():
                print(f"[DEBUG] tick[{key}] = {value}")
                logger.debug(f"[DEBUG] tick[{key}] = {value}")
                if 'dt' ==key:
                    tick[key] = pd.Timestamp.now()   # 缺啥补啥
            tick['dt'] = tick['dt'].strftime('%Y-%m-%d %H:%M:%S')
            tick['signal'] = do_t_signal(tick['code'], tick['price'])
            tick['gnn_score'] = gnn.get(tick['code'], 0)
            if tick['signal']['side'] in ('B', 'S'):
                order = Order(
                    code=tick['code'],
                    side=tick['signal']['side'],
                    price=tick['price'],
                    qty=tick['signal']['qty']
                )
                dry_run.save(order)
                asyncio.create_task(twap_worker(order.id))
            for ws in connected:
                await ws.send_text(json.dumps(tick))

        await asyncio.sleep(60)     # 每 60 s 一轮

connected = set()

@router.websocket("/ws/tick")
async def tick_ws(websocket: WebSocket):
    await websocket.accept()
    connected.add(websocket)
    try:
        while True:
            await websocket.receive_text()   # 保活
    except WebSocketDisconnect:
        connected.remove(websocket)
